# coding=utf-8

import torch
import torch.nn as nn
import torch.nn.functional as nn_functional
from rnn_net import getModel
from load_data import dataloader
import glob
from my_email import my_email

import csv

# 超参数
epoch_num = 300
lr = 1e-4

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# net
net = getModel().to(device)
# optimizer
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
# loss
loss_func = nn.CrossEntropyLoss()

log_path = '../csv/k.csv'
file = open(log_path, 'a+', encoding='utf-8', newline='')
csv_writer = csv.writer(file)
csv_writer.writerow([f'mode_name', 'Accuracy'])


def train():
    net.train()
    for epoch in range(epoch_num):

        for i, (x, y) in enumerate(dataloader):
            x = x.to(device)
            y = y.to(device)
            out = net(x)
            loss = loss_func(out, y)

            loss.backward()  # 根据loss计算模型的梯度
            optimizer.step()  # 根据梯度调整模型的参数
            optimizer.zero_grad()  # 梯度归0，进行下一轮计算

            acc = (out.argmax(dim=1) == y).sum().item() / len(y)
            print(
                'epoch:{},step:{}/{},loss:{},acc:{}'.format(epoch + 1, i + 1, len(dataloader), loss.item(), acc))

        torch.save(net, f'../model/net_epoch{epoch + 1}.pkl')


@torch.no_grad()  # 不计算模型梯度
def test_single(_net):
    net.eval()  # 进入测试模式

    acc = 0
    for i, (x, y) in enumerate(dataloader):
        x = x.to(device)
        y = y.to(device)
        out = _net(x).argmax(dim=1)
        acc += (out == y).sum().item()
    return acc


def test():
    modes_name = glob.glob('../model/net_epoch*.pkl')
    for mode_name in modes_name:
        new_net = torch.load(mode_name).to(device)
        acc = test_single(new_net)
        print(mode_name, acc / len(dataloader.dataset))
        csv_writer.writerow([mode_name, acc / len(dataloader.dataset)])
    file.close()


if __name__ == '__main__':
    print("开始训练啦")
    train()

    print('开始测试啦')
    test()
    print('发送邮件啦')
    my_email(__file__)
